{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Watersheds Segmentation <a href=\"https://mybinder.org/v2/gh/InsightSoftwareConsortium/SimpleITK-Notebooks/master?filepath=Python%2F32_Watersheds_Segmentation.ipynb\"><img style=\"float: right;\" src=\"https://mybinder.org/badge_logo.svg\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import SimpleITK as sitk\n",
    "from myshow import myshow, myshow3d\n",
    "\n",
    "# Download data to work on\n",
    "%run update_path_to_download_script\n",
    "from downloaddata import fetch_data as fdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = sitk.ReadImage(fdata(\"cthead1.png\"))\n",
    "myshow(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gradient Watersheds Segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = img.GetSpacing()[0]\n",
    "level = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_img = sitk.GradientMagnitude(img)\n",
    "myshow(feature_img, \"Edge Features\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws_img = sitk.MorphologicalWatershed(\n",
    "    feature_img, level=0, markWatershedLine=True, fullyConnected=False\n",
    ")\n",
    "myshow(sitk.LabelToRGB(ws_img), \"Watershed Over Segmentation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import interact, interactive, FloatSlider\n",
    "\n",
    "\n",
    "def callback(feature_img, *args, **kwargs):\n",
    "    ws_img = sitk.MorphologicalWatershed(feature_img, *args, **kwargs)\n",
    "\n",
    "    myshow(sitk.LabelToRGB(ws_img), \"Watershed Segmentation\")\n",
    "\n",
    "\n",
    "interact(\n",
    "    lambda **kwargs: callback(feature_img, **kwargs),\n",
    "    markWatershedLine=True,\n",
    "    fullyConnected=False,\n",
    "    level=FloatSlider(min=0, max=255, step=0.1, value=4.0),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Segmentation From Markers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_img = sitk.RegionalMinima(\n",
    "    feature_img,\n",
    "    backgroundValue=0,\n",
    "    foregroundValue=1.0,\n",
    "    fullyConnected=False,\n",
    "    flatIsMinima=True,\n",
    ")\n",
    "marker_img = sitk.ConnectedComponent(min_img)\n",
    "myshow(sitk.LabelToRGB(marker_img), \"Too many local minima markers\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws = sitk.MorphologicalWatershedFromMarkers(\n",
    "    feature_img, marker_img, markWatershedLine=True, fullyConnected=False\n",
    ")\n",
    "myshow(sitk.LabelToRGB(ws), \"Watershed Oversegmentation from markers\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt = [60, 60]\n",
    "idx = img.TransformPhysicalPointToIndex(pt)\n",
    "marker_img *= 0\n",
    "marker_img[0, 0] = 1\n",
    "marker_img[idx] = 2\n",
    "\n",
    "ws = sitk.MorphologicalWatershedFromMarkers(\n",
    "    feature_img, marker_img, markWatershedLine=True, fullyConnected=False\n",
    ")\n",
    "myshow(\n",
    "    sitk.LabelOverlay(img, ws, opacity=0.2),\n",
    "    \"Watershed Oversegmentation from manual markers\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Binary Watersheds for Object Separation "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb_img = sitk.ReadImage(fdata(\"coins.png\"))\n",
    "myshow(rgb_img, \"coins.png\")\n",
    "img = sitk.VectorIndexSelectionCast(rgb_img, 1)\n",
    "myshow(img, \"Green Coins\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_img = sitk.GradientMagnitudeRecursiveGaussian(img, sigma=1.5)\n",
    "myshow(feature_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws_img = sitk.MorphologicalWatershed(\n",
    "    feature_img, level=4, markWatershedLine=False, fullyConnected=False\n",
    ")\n",
    "myshow(sitk.LabelToRGB(ws_img), \"Watershed Over Segmentation\")\n",
    "seg = sitk.ConnectedComponent(ws_img != ws_img[0, 0])\n",
    "myshow(sitk.LabelOverlay(img, seg), \"Foreground Components\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filled = sitk.BinaryFillhole(seg != 0)\n",
    "d = sitk.SignedMaurerDistanceMap(\n",
    "    filled, insideIsPositive=False, squaredDistance=False, useImageSpacing=False\n",
    ")\n",
    "myshow(d, \"Inside Distance Map\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws = sitk.MorphologicalWatershed(d, markWatershedLine=False, level=1)\n",
    "myshow(sitk.LabelOverlay(img, ws))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws = sitk.Mask(ws, sitk.Cast(seg, sitk.sitkUInt8))\n",
    "myshow(sitk.LabelOverlay(img, ws), \"Split Objects\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-label Morphology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seg = ws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radius = 10\n",
    "bd_img = sitk.BinaryDilate(seg != 0, [radius] * seg.GetDimension())\n",
    "myshow(bd_img, \"Binary Dilate\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dist_img = sitk.SignedMaurerDistanceMap(\n",
    "    seg != 0, insideIsPositive=False, squaredDistance=False, useImageSpacing=False\n",
    ")\n",
    "wsd_img = sitk.MorphologicalWatershedFromMarkers(dist_img, seg, markWatershedLine=False)\n",
    "myshow(sitk.LabelOverlay(img, wsd_img))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "md_img = sitk.Mask(wsd_img, bd_img)\n",
    "myshow(sitk.LabelToRGB(md_img), \"Multi-label Dilate\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "e_img = sitk.BinaryErode(md_img != 0, [radius] * md_img.GetDimension())\n",
    "mo_img = sitk.Mask(md_img, e_img)\n",
    "myshow(sitk.LabelOverlay(img, mo_img), \"Multi-label Closing\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
